#include "mat_mul_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include "tiling/tiling_api.h"
using namespace matmul_tiling;


namespace optiling {
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{

//   MatMulTilingData tiling;
    //获取Tiling参数
    MatMulTilingData tiling;

    auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
    MultiCoreMatmulTiling cubeTiling(ascendcPlatform);  //创建Tiling对象

    auto shape_a = context->GetInputShape(0)->GetStorageShape();
    auto shape_b = context->GetInputShape(1)->GetStorageShape();
    // auto shape_bias = context->GetInputShape(2);
    
    int a_dim = shape_a.GetDimNum() - 2;
    int b_dim = shape_b.GetDimNum() - 2;
    int bias_dim = a_dim;

    int32_t M = shape_a.GetDim(a_dim);
    int32_t N = shape_b.GetDim(b_dim+1);
    int32_t K = shape_a.GetDim(a_dim+1);

    int a_num = 1;
    int b_num = 1;
    int bias_num = 1;
    
    int bias_shape_r[8];

    for(int i=0; i<a_dim; i++)
    {
        bias_shape_r[i] = shape_a.GetDim(i);
        a_num *= shape_a.GetDim(i);
    }
    for(int i=0; i<b_dim; i++)
    {
        b_num *= shape_b.GetDim(i);
    }
    bias_num = a_num;

    bool isBias = false;

    if (context->GetOptionalInputTensor(2) != nullptr)
    {
        isBias = true;
        auto shape_bias = context->GetInputShape(2)->GetStorageShape();
        bias_dim = shape_bias.GetDimNum() - 2;
        bias_num = 1;
        for(int i=0; i<bias_dim; i++)
        {
            bias_shape_r[i] = shape_bias.GetDim(i);
            bias_num *= bias_shape_r[i];
        }
    }

    
    int32_t c_dim = 1;
    int32_t c_num = 1;
    int32_t a_shape[8];
    int32_t b_shape[8];
    int32_t bias_shape[8];
    int32_t c_shape[8];

    c_dim = a_dim;
    if(c_dim < b_dim) c_dim = b_dim;
    if(c_dim < bias_dim) c_dim = bias_dim;


    uint32_t a_start = c_dim - a_dim;
    uint32_t b_start = c_dim - b_dim;
    uint32_t bias_start = c_dim - bias_dim;
    
    for(int i=0; i< c_dim; i++)
    {
        if(i < a_start) a_shape[i] = 1;
        else a_shape[i] = shape_a.GetDim(i - a_start);

        if(i < b_start) b_shape[i] = 1;
        else b_shape[i] = shape_b.GetDim(i - b_start);

        if(i < bias_start) bias_shape[i] = 1;
        else bias_shape[i] = bias_shape_r[i - bias_start];

        c_shape[i] = a_shape[i];
        if(c_shape[i] < b_shape[i]) c_shape[i] = b_shape[i];
        if(c_shape[i] < bias_shape[i]) c_shape[i] = bias_shape[i];

        c_num *= c_shape[i];
    }
    tiling.set_isBias(isBias);

    int32_t blockDim = 1;
    // int32_t c_big_num;
    int32_t c_small_num = c_num;
    int32_t big_num = c_num;
    // if(c_num >= 400)
    // {
        if(c_num <= 20)
        {
            blockDim = c_num;
            big_num = 0;
            // c_big_num = 1;
            c_small_num = 1;
        }
        else
        {
            blockDim = 20;
            c_small_num = c_num/20;
            big_num = c_num%20;
            if(big_num == 0)
            {
                big_num = 20;
                c_small_num -= 1;
            }
            // printf("!%d %d!",c_small_num, big_num);
        }
    // }

    // tiling.set_c_num(c_num);
    tiling.set_big_num(big_num);
    // tiling.set_c_big_num(c_big_num);
    tiling.set_c_small_num(c_small_num);
    tiling.set_c_dim(c_dim);

    tiling.set_a_shape(a_shape);
    tiling.set_b_shape(b_shape);
    tiling.set_bias_shape(bias_shape);
    tiling.set_c_shape(c_shape);

    cubeTiling.SetDim(1);
    cubeTiling.SetAType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);
    cubeTiling.SetBType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);
    cubeTiling.SetCType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);
    cubeTiling.SetBiasType(TPosition::GM, CubeFormat::ND, matmul_tiling::DataType::DT_FLOAT);
    cubeTiling.SetShape(M, N, K);
    cubeTiling.SetOrgShape(M, N, K);
    cubeTiling.SetBias(false);
    cubeTiling.SetBufferSpace(-1, -1, -1);
    cubeTiling.SetFixSplit(-1, -1, -1); 

    
    if (cubeTiling.GetTiling(tiling.cubeTilingData) == -1){ 
        return ge::GRAPH_FAILED;  
    }

    context->SetBlockDim(blockDim);

    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t userWorkspaceSize = blockDim*(M*K + K*N + M*N*2)*4*2;
    size_t systemWorkspaceSize = static_cast<size_t>(ascendcPlatform.GetLibApiWorkSpaceSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = userWorkspaceSize + systemWorkspaceSize;

  return ge::GRAPH_SUCCESS;
}
}


namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}


namespace ops {
class MatMul : public OpDef {
public:
    explicit MatMul(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_COMPLEX64})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_COMPLEX64})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Input("bias")
            .ParamType(OPTIONAL)
            .DataType({ge::DT_COMPLEX64})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});
        this->Output("z")
            .ParamType(REQUIRED)
            .DataType({ge::DT_COMPLEX64})
            .Format({ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND});

        this->SetInferShape(ge::InferShape);

        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend910b");

    }
};

OP_ADD(MatMul);
}
